#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: runyuanye@163.com
"""

import argparse
import os
import sys

sys.path.append('.')

"""
   对merge的人脸ID进行替换
"""


def read_file_list(file_list):
    faceid_subfaceid_files = []
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read File Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            # 015ff272-3768-4d4c-bebc-c14dc93065cd 477fc604-1479-4745-b8e9-9440a6ac8187 f8e4735b44724e25077106e0eaef15f0.jpg 0
            lines = line.strip().split()

            faceid = lines[0]
            subfaceid = lines[1]
            # file_name = lines[2]
            # mode = int(lines[3])

            faceid_subfaceid = '{} {}'.format(faceid, subfaceid)

            faceid_subfaceid_files.append([faceid_subfaceid, line])

        print('Read File Count: {}'.format(file_count))
    return faceid_subfaceid_files


def read_merge_file(file_list):
    faceid_subfaceids_map = {}
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read Feature Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            # NONE:960362bc-3dec-4785-8c87-bd60052bb7f0_192a5bc2-ef6f-4ef5-b82f-9bc6cd0791eb,960362bc-3dec-4785-8c87-bd60052bb7f0_29685424-4d56-49fb-acb7-3416fba1bb15,960362bc-3dec-4785-8c87-bd60052bb7f0_dc9d79c5-cccf-4e21-8ece-e2a14f00325e
            lines = line.strip().split(':')
            faceids = lines[1].split(',')

            faceid = faceids[0].replace('_', ' ', 1)
            for replace_faceid in faceids[1:]:
                faceid_subfaceids_map[replace_faceid.replace('_', ' ', 1)] = faceid
        print('Read Merge Count: {}'.format(file_count))
    return faceid_subfaceids_map


def main():
    parser = argparse.ArgumentParser(description="Replace Face ReID Image")
    parser.add_argument(
        "--file_list", default="/rootfs/media/kasim/Data/data/FaceImage/FaceRecSel.txt", help="path to image file list", type=str
    )
    parser.add_argument(
        "--merge_file", default="/rootfs/media/kasim/Data/data/FaceImage/MergeFace.txt", help="path to merge file", type=str
    )
    parser.add_argument(
        "--output_dir", default="/rootfs/media/kasim/Data/data/FaceImage", help="path to output", type=str
    )
    args = parser.parse_args()

    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        os.system('chmod a+wr {}'.format(output_dir))

    faceid_subfaceid_files = read_file_list(args.file_list)
    faceid_subfaceids_map = read_merge_file(args.merge_file)

    out_file_path = os.path.join(output_dir, 'Rep' + os.path.basename(args.file_list))
    with open(out_file_path, 'w') as out_file:
        for faceid_subfaceid, file in faceid_subfaceid_files:
            replace_faceid_subfaceid = faceid_subfaceids_map.get(faceid_subfaceid, None)
            if replace_faceid_subfaceid is not None:
                file = file.replace(faceid_subfaceid, replace_faceid_subfaceid, 1)
            out_file.write(file)
    os.system('chmod a+wr {}'.format(out_file_path))
    print('finish!')


if __name__ == '__main__':
    main()
